{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19f3b957",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "import os\n",
    "from PIL import Image\n",
    "from torch.utils import data\n",
    "from torchvision import transforms as T\n",
    "from sklearn.model_selection import train_test_split\n",
    "import torch.optim as optim\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.utils.data\n",
    "import torchvision.transforms as transforms\n",
    "from torch.autograd import Variable\n",
    "from torchvision.models import vgg16, VGG16_Weights\n",
    "import matplotlib.pyplot as plt\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b31cf580",
   "metadata": {},
   "source": [
    "PyTorch 深度学习项目的导入模块  \n",
    "编码声明：# -*- coding: utf-8 -*-    \n",
    "关于声明，自定义数据集为中文，为后续训练模型的准备，需要特意声明字符串类型，若这里不声明，Python会默认将字符串类型定义为ascii编码，导致训练模型时出现乱码。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81e0d7ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 定义标签\n",
    "# Labels = {'粑粑柑': 0, '白兰瓜': 1, '白萝卜': 2, '白心火龙果': 3, '百香果': 4, '菠萝': 5, '菠萝莓': 6, '菠萝蜜': 7, '草莓': 8, '车厘子': 9, '番石榴-百': 10, '番石榴-红': 11, '佛手瓜': 12, '甘蔗': 13, '桂圆': 14, '哈密瓜': 15, '黑莓': 16, '红苹果': 17, '红心火龙果': 18, '胡萝卜': 19, '黄桃': 20, '金桔': 21, '橘子': 22, '蓝莓': 23, '梨': 24, '李子': 25, '荔枝': 26, '莲雾': 27, '榴莲': 28, '芦柑': 29, '芒果': 30, '毛丹': 31, '猕猴桃': 32, '木瓜': 33, '柠檬': 34, '牛油果': 35, '蟠桃': 36, '枇杷': 37, '葡萄-白': 38, '葡萄-红': 39, '脐橙': 40, '青柠': 41, '青苹果': 42, '人参果': 43, '桑葚': 44, '沙果': 45, '沙棘': 46, '砂糖橘': 47, '山楂': 48, '山竹': 49, '蛇皮果': 50, '圣女果': 51, '石 '人参果': '人 '人参果': 43, '桑葚': 44, '沙果': 45, '沙棘': 46, '砂糖橘': 47, '山楂': 48, '山竹': 49, '蛇皮果': 50, '圣女果': 51, '石榴': 52, '柿子': 53, '树莓': 54, '水蜜桃': 55, '酸角': 56, '甜瓜-白': 57, '甜瓜-金': 58, '甜瓜-绿': 59, '甜瓜-伊丽莎白': 60, '沃柑': 61, '无花果': 62, '西瓜': 63, '西红柿': 64, '西梅': 65, '西柚': 66, '香蕉': 67, '香橼': 68, '杏': 69, '血橙': 70, '羊角蜜': 71, '羊奶果': 72, '杨梅': 73, '杨桃': 74, '腰果': 75, '椰子': 76, '樱桃': 77, '油桃': 78, '柚子': 79, '枣': 80}\n",
    "# 自定义数据集类\n",
    "Labels = {\n",
    "    '粑粑柑': 0, '白兰瓜': 1, '白萝卜': 2, '白心火龙果': 3, '百香果': 4, '菠萝': 5, '菠萝莓': 6, '菠萝蜜': 7, '草莓': 8, '车厘子': 9, \n",
    "    '番石榴-百': 10, '番石榴-红': 11, '佛手瓜': 12, '甘蔗': 13, '桂圆': 14, '哈密瓜': 15, '黑莓': 16, '红苹果': 17, '红心火龙果': 18, \n",
    "    '胡萝卜': 19, '黄桃': 20, '金桔': 21, '橘子': 22, '蓝莓': 23, '梨': 24, '李子': 25, '荔枝': 26, '莲雾': 27, '榴莲': 28, '芦柑': 29, \n",
    "    '芒果': 30, '毛丹': 31, '猕猴桃': 32, '木瓜': 33, '柠檬': 34, '牛油果': 35, '蟠桃': 36, '枇杷': 37, '葡萄-白': 38, '葡萄-红': 39, \n",
    "    '脐橙': 40, '青柠': 41, '青苹果': 42, '人参果': 43, '桑葚': 44, '沙果': 45, '沙棘': 46, '砂糖橘': 47, '山楂': 48, '山竹': 49, \n",
    "    '蛇皮果': 50, '圣女果': 51, '石榴': 52, '柿子': 53, '树莓': 54, '水蜜桃': 55, '酸角': 56, '甜瓜-白': 57, '甜瓜-金': 58, '甜瓜-绿': 59, \n",
    "    '甜瓜-伊丽莎白': 60, '沃柑': 61, '无花果': 62, '西瓜': 63, '西红柿': 64, '西梅': 65, '西柚': 66, '香蕉': 67, '香橼': 68, '杏': 69, \n",
    "    '血橙': 70, '羊角蜜': 71, '羊奶果': 72, '杨梅': 73, '杨桃': 74, '腰果': 75, '椰子': 76, '樱桃': 77, '油桃': 78, '柚子': 79, '枣': 80\n",
    "}\n",
    "\n",
    "class SeedlingData(data.Dataset):\n",
    "    def __init__(self, root, transforms=None, train=True, test=False):\n",
    "        self.test = test\n",
    "        self.transforms = transforms\n",
    "        # 加载图像路径\n",
    "        if self.test:\n",
    "            self.imgs = [os.path.join(root, img) for img in os.listdir(root)]\n",
    "        else:\n",
    "            imgs_labels = [os.path.join(root, img) for img in os.listdir(root)]\n",
    "            imgs = []\n",
    "            for imglable in imgs_labels:\n",
    "                for imgname in os.listdir(imglable):\n",
    "                    imgpath = os.path.join(imglable, imgname)\n",
    "                    imgs.append(imgpath)\n",
    "            trainval_files, val_files = train_test_split(imgs, test_size=0.3, random_state=42)\n",
    "            self.imgs = trainval_files if train else val_files\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        img_path = self.imgs[index]\n",
    "        if self.test:\n",
    "            label = -1\n",
    "        else:\n",
    "            labelname = os.path.basename(os.path.dirname(img_path))\n",
    "            label = Labels[labelname]\n",
    "        data = Image.open(img_path).convert('RGB')\n",
    "        data = self.transforms(data)\n",
    "        return data, label\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.imgs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b402d71",
   "metadata": {},
   "source": [
    "这段代码定义了一个用于水果图像分类的数据集处理模块，主要包含以下内容：  \n",
    "1. 标签映射定义  \n",
    "标签字典 Labels：将 81 种水果名称映射到 0-80 的数字标签。为utf-8编码格式。  \n",
    "标签覆盖了常见水果及部分变种（如番石榴 - 白、番石榴 - 红），结构清晰便于模型输出结果映射。  \n",
    "2. 自定义数据集类 SeedlingData  \n",
    "初始化方法 __init__  \n",
    "参数说明：  \n",
    "root：数据集根目录路径  \n",
    "transforms：图像预处理转换函数  \n",
    "train：是否为训练集（默认True）  \n",
    "test：是否为测试集（默认False）  \n",
    "数据加载逻辑：  \n",
    "测试集：直接加载根目录下所有图像  \n",
    "训练 / 验证集：  \n",
    "遍历根目录下的所有标签文件夹（每个文件夹名对应一个水果类别）  \n",
    "收集所有图像路径  \n",
    "使用train_test_split按 7:3 比例划分训练集和验证集（随机种子 42 保证可复现）  \n",
    "数据获取方法 __getitem__  \n",
    "输入处理：  \n",
    "根据索引获取图像路径  \n",
    "读取图像并转换为 RGB 格式  \n",
    "应用预处理转换（如缩放、归一化）  \n",
    "标签处理：  \n",
    "测试集标签设为 - 1（无真实标签）  \n",
    "训练 / 验证集通过文件夹名从Labels字典获取对应标签  \n",
    "返回值：预处理后的图像张量和标签（训练 / 验证集）或仅图像（测试集）  \n",
    "长度获取方法 __len__  \n",
    "返回数据集包含的图像总数，用于数据加载器批量处理。  \n",
    "3. 核心功能总结  \n",
    "数据组织：支持按标签文件夹结构组织数据集（root/类别名/图像.jpg）  \n",
    "数据集划分：自动将数据划分为训练集（70%）和验证集（30%）  \n",
    "预处理支持：通过transforms参数灵活配置图像转换流程  \n",
    "多模式支持：可处理训练集、验证集和测试集三种场景  \n",
    "4. 应用场景  \n",
    "该数据集类适用于基于 PyTorch 的水果图像分类任务，可与之前的 VGG16 模型结合使用，实现：  \n",
    "\n",
    "训练阶段：加载训练集数据并优化模型参数  \n",
    "验证阶段：评估模型在验证集上的性能   \n",
    "测试阶段：对未知图像进行预测（仅加载图像，无标签） \n",
    "\n",
    "通过这种设计，代码实现了数据加载、预处理和标签映射的自动化，便于后续模型训练与部署  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d51aebb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 超参数\n",
    "modellr = 1e-4\n",
    "BATCH_SIZE = 32\n",
    "EPOCHS = 10\n",
    "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# 数据预处理\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((224, 224)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # ImageNet 标准化参数\n",
    "])\n",
    "\n",
    "# 加载数据集\n",
    "train_dataset = SeedlingData(root='best.1/fruit81_full', transforms=transform, train=True)\n",
    "val_dataset = SeedlingData(root='best.1/fruit81_full', transforms=transform, train=False)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)\n",
    "\n",
    "# 加载预训练的 VGG16 模型\n",
    "model_ft = vgg16(weights=VGG16_Weights.IMAGENET1K_V1)\n",
    "num_ftrs = model_ft.classifier[6].in_features\n",
    "model_ft.classifier[6] = nn.Linear(num_ftrs, len(Labels))\n",
    "model_ft = model_ft.to(DEVICE)\n",
    "\n",
    "# 定义优化器和损失函数\n",
    "optimizer = optim.Adam(model_ft.parameters(), lr=modellr)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "# 学习率调整函数\n",
    "def adjust_learning_rate(optimizer, epoch):\n",
    "    modellrnew = modellr * (0.1 ** (epoch // 50))\n",
    "    for param_group in optimizer.param_groups:\n",
    "        param_group['lr'] = modellrnew\n",
    "\n",
    "# 训练函数\n",
    "def train(model, device, train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    sum_loss = 0\n",
    "    correct = 0\n",
    "    total_num = len(train_loader.dataset)\n",
    "\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        output = model(data)\n",
    "        loss = criterion(output, target)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        sum_loss += loss.item()\n",
    "        _, pred = torch.max(output, 1)\n",
    "        correct += pred.eq(target).sum().item()\n",
    "\n",
    "        if (batch_idx + 1) % 10 == 0:\n",
    "            print(f'Train Epoch: {epoch} [{(batch_idx + 1) * len(data)}/{total_num} ({100. * (batch_idx + 1) / len(train_loader):.0f}%)]\\tLoss: {loss.item():.6f}')\n",
    "\n",
    "    avg_loss = sum_loss / len(train_loader)\n",
    "    accuracy = 100. * correct / total_num\n",
    "    print(f'Epoch: {epoch}, Loss: {avg_loss:.6f}, Accuracy: {accuracy:.2f}%')\n",
    "\n",
    "# 验证函数\n",
    "def val(model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    total_num = len(test_loader.dataset)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            loss = criterion(output, target)\n",
    "            test_loss += loss.item()\n",
    "            _, pred = torch.max(output, 1)\n",
    "            correct += pred.eq(target).sum().item()\n",
    "\n",
    "    avg_loss = test_loss / len(test_loader)\n",
    "    accuracy = 100. * correct / total_num\n",
    "    print(f'\\nValidation set: Average loss: {avg_loss:.4f}, Accuracy: {correct}/{total_num} ({accuracy:.2f}%)\\n')\n",
    "\n",
    "# 训练和验证\n",
    "for epoch in range(1, EPOCHS + 1):\n",
    "    adjust_learning_rate(optimizer, epoch)\n",
    "    train(model_ft, DEVICE, train_loader, optimizer, epoch)\n",
    "    val(model_ft, DEVICE, val_loader)\n",
    "\n",
    "# 保存模型\n",
    "torch.save(model_ft.state_dict(), 'model.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c705474",
   "metadata": {},
   "source": [
    "这段代码实现了一个基于 VGG16 预训练模型的水果图像分类系统的训练流程，主要包含以下核心模块：  \n",
    "1. 超参数与数据预处理  \n",
    "超参数设置：  \n",
    "学习率 modellr=1e-4，批量大小 BATCH_SIZE=32，训练轮次 EPOCHS=10  \n",
    "设备自动检测：优先使用 GPU（cuda），否则用 CPU  \n",
    "图像预处理：  \n",
    "Resize 至 224×224 像素，转为 Tensor 格式  \n",
    "使用 ImageNet 标准的归一化参数（均值[0.485, 0.456, 0.406]，标准差[0.229, 0.224, 0.225]）  \n",
    "2. 数据集加载与划分  \n",
    "使用自定义数据集类SeedlingData：  \n",
    "训练集：从fruit81_full文件夹加载数据，自动按 7:3 比例划分训练集和验证集  \n",
    "验证集：通过train=False参数加载剩余 30% 数据  \n",
    "数据加载器：  \n",
    "训练集：打乱顺序（shuffle=True），便于模型学习  \n",
    "验证集：不打乱（shuffle=False），保证评估一致性  \n",
    "3. 模型构建与优化配置  \n",
    "迁移学习模型：  \n",
    "加载 VGG16 预训练模型（IMAGENET1K_V1权重）  \n",
    "修改分类层：将原 1000 维输出改为 81 维（对应 81 种水果标签）  \n",
    "优化器与损失函数：  \n",
    "优化器：Adam 算法（自适应学习率）  \n",
    "损失函数：交叉熵损失（适用于多分类任务）  \n",
    "4. 学习率调整与训练逻辑  \n",
    "学习率衰减策略：  \n",
    "每 50 轮次学习率乘以 0.1（0.1 ** (epoch // 50)）  \n",
    "训练函数train：  \n",
    "模型设为训练模式（model.train()）  \n",
    "批量处理数据，计算损失并反向传播  \n",
    "每 10 个批次打印进度， epoch 结束后输出平均损失和准确率  \n",
    "验证函数val：  \n",
    "模型设为评估模式（model.eval()）  \n",
    "禁用梯度计算，计算验证集损失和准确率  \n",
    "5. 模型训练与保存  \n",
    "循环训练流程：  \n",
    "对每个 epoch，先调整学习率，再执行训练和验证  \n",
    "模型保存：  \n",
    "训练完成后保存模型参数到model.pth文件  \n",
    "核心技术点总结  \n",
    "迁移学习：利用 VGG16 在 ImageNet 上的预训练权重，加速水果分类模型的收敛  \n",
    "数据处理：通过自定义数据集类实现标签映射和自动划分，保证数据加载效率  \n",
    "优化策略：Adam 优化器结合学习率衰减，平衡训练速度和收敛效果  \n",
    "工程化设计：分离训练、验证逻辑，便于调试和性能监控  \n",
    "应用场景  \n",
    "该代码可用于构建一个包含 81 类水果的图像分类系统，训练后的模型可集成到 PyQt 界面中（如之前的FruitRecognitionApp），实现摄像头实时识别或图片上传识别功能。如需提升精度，可增加训练轮次、调整数据增强策略或更换更复杂的 backbone 模型。"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
